package com.binance.graceful.shutdown.core;

import lombok.extern.log4j.Log4j2;
import org.apache.catalina.connector.Connector;
import org.springframework.boot.context.embedded.tomcat.TomcatConnectorCustomizer;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextClosedEvent;

import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * Created by robin.wu on 2018/10/26.
 */
@Log4j2
public class GracefulShutdownTomcat implements TomcatConnectorCustomizer, ApplicationListener<ContextClosedEvent> {

    private volatile Connector connector;
    private Long waitTime;    //default time to wait some task not finished(ms)
    private Long accessTime;  //time to access when kill signal


    public GracefulShutdownTomcat(Long waitTime) {
        this(waitTime,0L);
    }

    public GracefulShutdownTomcat(Long waitTime,Long accessTime) {
        if(waitTime == null) {
            this.waitTime = 30000L;
        }else {
            this.waitTime = waitTime;
        }

        if(accessTime == null) {
            this.accessTime = 30000L;
        }else {
            this.accessTime = accessTime;
        }
    }

    @Override
    public void customize(Connector connector) {
        this.connector = connector;
    }

    @Override
    public void onApplicationEvent(ContextClosedEvent event) {
        shouldSleep();

        this.connector.pause();
        Executor executor = this.connector.getProtocolHandler().getExecutor();
        if (executor instanceof ThreadPoolExecutor) {
            try {
                ThreadPoolExecutor threadPoolExecutor = (ThreadPoolExecutor) executor;
                threadPoolExecutor.shutdown();
                if (!threadPoolExecutor.awaitTermination(waitTime, TimeUnit.MILLISECONDS)) {
                    log.warn("Tomcat thread pool did not shut down gracefully within " + waitTime + " seconds. Proceeding with forceful shutdown");
                }
            } catch (InterruptedException ex) {
                Thread.currentThread().interrupt();
            }
        }
    }

    private void shouldSleep(){
        if(this.accessTime >0L) {
            try {
                Thread.sleep(accessTime);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                log.error("GracefulShutdownTomcat error",e);
            }
        }
    }
}
